import numpy as np
import matplotlib.pyplot as plt


# model
def model(XX, theta):
    return XX.dot(theta)


# sigmoid function
def sigmoid(z):
    return 1 / (1 + np.exp(-z))


# cost function with regularization
def cost_func(h, y, lam=0, theta=None):
    m = len(h)
    if lam != 0:
        theta_cp = theta.copy()
        theta_cp[0] = 0
        r = 2 / m * np.sum(theta_cp ** 2)
    return -1 / m * np.sum(y * np.log(h) + (1 - y) * np.log(1 - h))


# gradient descent with regularization
def grad(XX, y, lam=0, alpha=0.001, iter0=15000):
    m, n = XX.shape
    group = iter0 // 20
    theta = np.zeros(n)
    j_his = np.zeros(iter0)
    for i in range(iter0):
        z = model(XX, theta)
        h = sigmoid(z)
        j = cost_func(h, y, lam, theta)
        if i % group == 0:
            print(f'#{i + 1} cost func value = {j}')
        j_his[i] = j
        dt = 1 / m * XX.T.dot(h - y)
        if lam != 0:
            theta_cp = theta.copy()
            theta_cp[0] = 0
            r = lam / m * theta_cp
            dt += r
        theta -= alpha * dt
    if i % group != 0:
        print(f'#{i + 1} cost func value = {j}')
    return theta, j_his, h


# accuracy function
def score(h, y):
    return np.mean(y == (h > 0.5))


# main program
if '__main__' == __name__:
    plt.figure(figsize=[12, 8])
    spr = 2
    spc = 2
    spn = 0

    # plot sigmoid function
    spn += 1
    plt.subplot(spr, spc, spn)
    plt_x = np.linspace(-10, 10, 1000)
    plt_y = sigmoid(plt_x)
    plt.plot(plt_x, plt_y, label='sigmoid')
    plt.grid()
    plt.legend()

    # load data
    data_train = np.loadtxt('d1.txt', delimiter=',')
    m_train = len(data_train)
    x_train = data_train[:, :-1]
    y_train = data_train[:, -1]
    data_test = np.loadtxt('d2.txt', delimiter=',')
    m_test = len(data_test)
    x_test = data_test[:, :-1]
    y_test = data_test[:, -1]

    # scale data
    mu = x_train.mean(axis=0)
    sigma = x_train.std(axis=0)
    x_train -= mu
    x_train /= sigma
    mu = x_test.mean(axis=0)
    sigma = x_test.std(axis=0)
    x_test -= mu
    x_test /= sigma

    # splice data
    XX_train = np.c_[np.ones(m_train), x_train]
    XX_test = np.c_[np.ones(m_test), x_test]
    # np.savetxt('trainXX.tmp.txt', XX_train)  # tmp
    # np.savetxt('testXX.tmp.txt', XX_test)  # tmp

    # fit training data
    lam = 3
    alpha = 0.001
    iter0 = 15000
    theta, j_his, h_train = grad(XX_train, y_train, lam, alpha, iter0)
    print(f'Theta = {theta}')
    print(f'Training score = {score(h_train, y_train)}')

    # plot cost function values
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.plot(j_his, label='cost function')
    plt.xlabel('iterations')
    plt.grid()
    plt.legend()

    # test on testing data
    z_test = model(XX_test, theta)
    h_test = sigmoid(z_test)
    print(f'Testing score = {score(h_test, y_test)}')

    # translate theta for plotting border line
    def get_border(x, theta):
        return - (theta[0] + theta[2] * x) / theta[3]

    # plot X2, X3 on training data
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Training data')
    pos_idx = y_train == 1
    neg_idx = np.invert(pos_idx)
    plt.scatter(x_train[pos_idx, 1], x_train[pos_idx, 2], c='y', s=1, label='positive')
    plt.scatter(x_train[neg_idx, 1], x_train[neg_idx, 2], c='g', s=1, label='negative')
    plt_line_x = np.array([x_train[:, 1].min(), x_train[:, 1].max()])
    plt_line_y = get_border(plt_line_x, theta)
    plt.plot(plt_line_x, plt_line_y, 'r-', label='border')
    plt.xlabel('X2')
    plt.ylabel('X3')
    plt.xlim([-1.5, 1.5])
    plt.ylim([-3, 3])
    plt.grid()
    plt.legend()

    # plot X2, X3 on testing data
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Testing data')
    pos_idx = y_test == 1
    neg_idx = np.invert(pos_idx)
    plt.scatter(x_test[pos_idx, 1], x_test[pos_idx, 2], c='y', s=1, label='positive')
    plt.scatter(x_test[neg_idx, 1], x_test[neg_idx, 2], c='g', s=1, label='negative')
    plt_line_x = np.array([x_test[:, 1].min(), x_test[:, 1].max()])
    plt_line_y = get_border(plt_line_x, theta)
    plt.plot(plt_line_x, plt_line_y, 'r-', label='border')
    plt.xlabel('X2')
    plt.ylabel('X3')
    plt.xlim([-1.5, 1.5])
    plt.ylim([-3, 3])
    plt.grid()
    plt.legend()

    # lastly, show all drawings
    plt.show()
